#include "solver_stable.h"

#include <math.h>
#include <memory.h>
#include <omp.h>

double* solve_stable(double* u0, Area area, double epsilon, TwoDimFunc f)
{
	int n = area.n;
	double a = area.a, b = area.b;
	double* u_cur = new double[n * n];
	double* u_prev = new double[n * n];
	double h = (b - a) / n;
	double error;
	double start, finish;

	double* f_cur = new double[n * n];

	start = PortableGetTime();

#pragma omp parallel for shared( n, f_cur, f, h) 
	for (int i = 0; i < n; ++i)
	{
#pragma omp parallel for 
		for (int j = 0; j < n; ++j)
		{
			f_cur[i * n + j] = f(a + i * h, a + j * h);
		}
	}

	finish = PortableGetTime();
	printf("Compute function time: %.4f\n", finish - start);

#pragma omp sections
	{
#pragma omp section 
		memcpy(u_prev, u0, n * n * sizeof(double));
#pragma omp section 
		memcpy(u_cur, u0, n * n * sizeof(double));
	}

	start = PortableGetTime();
	do
	{ 

#pragma omp parallel for
		for (int i = 1; i < n - 1; ++i)
		{
			for (int j = 1; j < n - 1; ++j)
			{
				u_cur[i * n + j] = 0.25 * 
					(u_prev[(i - 1) * n + j] + u_prev[(i + 1) * n + j] +
					u_prev[i * n + j - 1] + u_prev[i * n + j + 1] - 
					h * h * f_cur[i * n + j]);
			}
		}

		double e, *err;
		error = abs(u_cur[0] - u_prev[0]);

		int countSteams = omp_get_max_threads();

		int col = (n * n - 2)/countSteams;

		err = new double[countSteams];

		for(int i = 0; i < countSteams; i++)
			err[i] = error;

#pragma omp parallel private(e)
		{
			int ID = omp_get_thread_num();

			for (int i = 1; i <= col ; ++i)
			{
				e = abs(u_cur[ID*col + i] - u_prev[ID*col + i]);
				if (e > err[ID]) err[ID] = e;
			}

		}

		error = err[0];
		for(int i = 1; i < countSteams; i++)
			if(error < err[i])
				error = err[i];

		delete[] err;

		double* tmp;
		tmp = u_cur; u_cur = u_prev; u_prev = tmp;

	} while (error > epsilon);
	finish = PortableGetTime();

	printf("Solve time: %.4f\n", finish - start);

	delete[] u_cur;
	return u_prev;
}